package _13_red_black_tree.case1;

import _00_my.annotation.Risky;

import static _13_red_black_tree.case1.RedBlackTree.Color.BLACK;
import static _13_red_black_tree.case1.RedBlackTree.Color.RED;

/**
 * @author: mornd
 * @dateTime: 2023/6/15 - 19:52
 * 红黑树定义
 */
@Risky
class RedBlackTree<E> {
    RBNode root;

    /**
     * 新增
     * 遇到红红不平衡（4种情况）进行调整，否则正常新增
     *
     * @param key
     * @param value
     */
    public void put(int key, E value) {
        RBNode curr = root;
        // 父节点
        RBNode parent = null;
        while (curr != null) {
            parent = curr;
            if (key < curr.key) {
                curr = curr.left;
            } else if (key > curr.key) {
                curr = curr.right;
            } else {
                // 已找到，更新 value
                curr.value = value;
                return;
            }
        }

        RBNode inserted = new RBNode(key, value);
        if (parent == null) {
            root = inserted;
        } else if (key < parent.key) {
            // 将 parent 的左节点更新为插入节点
            parent.left = inserted;
            inserted.parent = parent;
        } else {
            parent.right = inserted;
            inserted.parent = parent;
        }

        fixRedRed(inserted);

//        System.out.println(this + "\n-----------------------------------------");
    }

    /**
     * 遇到红红问题进行修复调整
     *
     * @param inserted
     */
    private void fixRedRed(RBNode inserted) {
        /*
            红红的四种情况
                插入节点视为红色(这里插入节点默认是红色)
                    case1：插入节点为根节点，直接将根节点变为黑色即可
                    case2：插入节点的父亲若为黑色，树的红黑性质不变，不需要调整
                插入节点的父亲为红色，存放红红相邻
                    case3：叔叔为红色
                    case4：叔叔为黑色
         */
        if (root == inserted) {
            // case1 根变为黑色即可
            inserted.color = BLACK;
            return;
        }

        if (isBlack(inserted.parent)) {
            // case2 红黑不处理
            return;
        }

        // 代码走到这里说明，插入节点的父亲也是红色

        RBNode parent = inserted.parent;
        RBNode uncle = inserted.uncle();
        // 祖父节点
        RBNode grandparent = parent.parent;

        if (isRed(uncle)) {
            /*
            case3
            叔叔为红色处理：inserted 的父亲和叔叔都变为黑色，祖父变为红色，最后递归向上调用
             */
            parent.color = BLACK;
            uncle.color = BLACK;
            grandparent.color = RED;

            fixRedRed(grandparent);
        } else {
            /*
            case4
            叔叔为黑色处理：
                 此时又分4种情况 LL LR RL RR
                    变色+旋转
             */
            if (parent.isLeftChild() && inserted.isLeftChild()) {
                // 父亲为左孩子，插入节点也为左孩子，即左左失衡，需要变色+右旋 LL
                parent.color = BLACK;
                grandparent.color = RED;
                rightRotate(grandparent);
            } else if (parent.isLeftChild() && !inserted.isLeftChild()) {
                // 父亲为左孩子，插入节点也为右孩子，即左右失衡，需要变色+右旋 // LR
                leftRotate(parent);
                inserted.color = BLACK;
                grandparent.color = RED;
                rightRotate(grandparent);
            } else if (!parent.isLeftChild() && !inserted.isLeftChild()) {
                // RR
                parent.color = BLACK;
                grandparent.color = RED;
                leftRotate(grandparent);
            } else {
                // RL
                rightRotate(parent);
                inserted.color = BLACK;
                grandparent.color = RED;
                leftRotate(grandparent);
            }
        }
    }

    /**
     * 删除
     * 遇到黑黑不平衡（6种情况）进行调整，否则正常删除
     *
     * @param key
     */
    public void remove(int key) {
        RBNode deleted = find(key);
        if (deleted == null) {
            return;
        }
        remove(deleted);
    }

    private void remove(RBNode deleted) {
        RBNode replaced = findReplaced(deleted);
        RBNode parent = deleted.parent;

        if (replaced == null) {
            // 没有孩子
            if (root == deleted) {
                // case1：被删除节点是 root
                root = null;
            } else {
                // case2
                if (isBlack(deleted)) {
                    fixDoubleBlack(deleted);
                } else {
                    // nothing...
                }
                if (deleted.isLeftChild()) {
                    parent.left = null;
                } else {
                    parent.right = null;
                }
                deleted.parent = null;
            }
            return;
        }
        if (replaced.left == null || replaced.right == null) {
            // 有一个孩子
            if (root == deleted) {
                // case1
                root.key = replaced.key;
                root.value = replaced.value;
                root.left = root.right = null;
            } else {
                // case2
                if (deleted.isLeftChild()) {
                    parent.left = replaced;
                } else {
                    parent.right = replaced;
                }
                replaced.parent = parent;
                deleted.left = deleted.right = deleted.parent = null; // help GC
                if (isBlack(deleted) && isBlack(replaced)) {
                    fixDoubleBlack(replaced);
                } else {
                    replaced.color = BLACK;
                }
            }
            return;
        }
        // case0：有两个孩子

        // 交换待删除节点和其后继节点的 key 和 value
        int k = deleted.key;
        deleted.key = replaced.key;
        replaced.key = k;

        E v = deleted.value;
        deleted.value = replaced.value;
        replaced.value = v;

        // 此时不用删除待删除节点，而是删除后继节点，因为后继节点替死了，李代桃僵
        remove(replaced);
    }

    /**
     * 处理双黑（删除的节点和剩下的节点都是黑色）
     *
     * @param n
     */
    private void fixDoubleBlack(RBNode n) {
        if (n == root) {
            return;
        }
        RBNode parent = n.parent;
        RBNode sibling = n.sibling();
        if (isRed(sibling)) {
            if (n.isLeftChild()) {
                leftRotate(parent);
            } else {
                rightRotate(parent);
            }
            parent.color = RED;
            sibling.color = BLACK;
            fixDoubleBlack(n);
            return;
        }

        if (sibling != null) {
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                sibling.color = RED;
                if (isRed(parent)) {
                    parent.color = BLACK;
                } else {
                    fixDoubleBlack(parent);
                }
            } else {
                if (sibling.isLeftChild() && isRed(sibling.left)) {
                    // LL
                    rightRotate(parent);
                    sibling.left.color = BLACK;
                    sibling.color = parent.color;
                } else if (sibling.isLeftChild() && isRed(sibling.right)) {
                    // LR
                    sibling.right.color = parent.color;
                    leftRotate(sibling);
                    rightRotate(parent);
                } else if (!sibling.isLeftChild() && isRed(sibling.left)) {
                    // RL
                    sibling.left.color = parent.color;
                    rightRotate(sibling);
                    leftRotate(parent);
                } else {
                    // RR
                    leftRotate(parent);
                    sibling.right.color = BLACK;
                    sibling.color = parent.color;
                }
                parent.color = BLACK;
            }
        } else {
            fixDoubleBlack(parent);
        }
    }

    /**
     * 根据 key 查找节点
     */
    private RBNode find(int k) {
        RBNode curr = root;
        while (curr != null) {
            if (k < curr.key) {
                curr = curr.left;
            } else if (k > curr.key) {
                curr = curr.right;
            } else {
                return curr;
            }
        }
        return null;
    }

    /**
     * 查找被删除节点的剩余节点
     *
     * @param deleted
     * @return
     */
    private RBNode findReplaced(RBNode deleted) {
        if (deleted.left == null && deleted.right == null) {
            return null;
        }
        if (deleted.left == null) {
            return deleted.right;
        }
        if (deleted.right == null) {
            return deleted.left;
        }

        // 左右都不为空，找到后继节点并返回
        RBNode s = deleted.right;
        while (s.left != null) {
            s = s.left;
        }
        return s;
    }

    /**
     * 根据 key 查找 value
     */
    public E get(int k) {
        RBNode n = find(k);
        return n == null ? null : n.value;
    }

    /**
     * 左旋
     *
     * @param pink
     */
    private void leftRotate(RBNode pink) {
        RBNode parent = pink.parent;

        RBNode yellow = pink.right;
        RBNode green = yellow.left;
        yellow.left = pink;
        pink.right = green;

        yellow.parent = parent;
        pink.parent = yellow;
        if (green != null) {
            green.parent = pink;
        }


        if (parent != null) {
            // 建立父节点与左旋节点的关系
            if (parent.left == pink) {
                parent.left = yellow;
            } else {
                parent.right = yellow;
            }
        } else {
            root = yellow;
        }
    }

    // 右旋
    private void rightRotate(RBNode pink) {
        RBNode parent = pink.parent;
        // 开始右旋，参照图片 img/右旋前.png
        RBNode yellow = pink.left;
        RBNode green = yellow.right;
        yellow.right = pink;
        pink.left = green;

        // 3个节点全部都要换爹
        yellow.parent = parent;
        pink.parent = yellow;
        if (green != null) {
            green.parent = pink;
        }

        if (parent != null) {
            // 建立父节点与右旋节点的关系
            if (parent.left == pink) {
                parent.left = yellow;
            } else {
                parent.right = yellow;
            }
        } else {
            root = yellow;
        }
    }

    // 判断是否是红色 如果节点为 null，则为黑色
    private boolean isRed(RBNode n) {
        return n != null && n.color == RED;
    }

    // 判断是否是黑色
    private boolean isBlack(RBNode n) {
        return !isRed(n);
    }

    class RBNode {
        int key;
        E value;
        RBNode left;
        RBNode right;
        RBNode parent;
        // 节点颜色，默认红色
        Color color = RED;

        // 判断该节点是否是左孩子
        boolean isLeftChild() {
            return parent != null && parent.left == this;
        }

        // 获取该节点的叔叔节点，并且最多只能有一个
        RBNode uncle() {
            if (parent == null || parent.parent == null) {
                return null;
            }
            if (parent.isLeftChild()) {
                return parent.parent.right;
            } else {
                return parent.parent.left;
            }
        }

        // 获取兄弟节点
        RBNode sibling() {
            if (parent == null) {
                return null;
            }
            if (isLeftChild()) {
                return parent.right;
            } else {
                return parent.left;
            }
        }

        public RBNode(int key, E value) {
            this.key = key;
            this.value = value;
        }

        public RBNode(int key, E value, RBNode parent, Color color) {
            this.key = key;
            this.value = value;
            this.parent = parent;
            this.color = color;
        }

        public RBNode(int key, E value, RBNode parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }
    }

    enum Color {
        RED,
        BLACK;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        printVerticalTree(root, "", sb);
        return sb.toString();
    }

    private void printVerticalTree(RBNode node, String prefix, StringBuilder sb) {
        if (node == null) {
            return;
        }
        printVerticalTree(node.right, prefix + "        ", sb);
        sb.append(node.color == RED ? "\u001B[31m" : "\u001B[30m");
        sb.append(prefix).append("|--");
        sb.append(node.key).append("\u001B[0m").append("\n");
        printVerticalTree(node.left, prefix + "        ", sb);
    }
}
